Models Documentation

This documentation provides an overview of key models, focusing on essential input parameters and high-level functionality.

Conditional Entropy Models

class encoding_information.models.conditional_entropy_models.AnalyticGaussianNoiseModel(sigma)

Bases: MeasurementNoiseModel

Analytical model for estimating the conditional entropy H(Y | X) when the noise proces is additive independent Gaussian noise at each pixel.

Initialize the AnalyticGaussianNoiseModel.

Parameters:

sigma (float) – Standard deviation of the Gaussian noise.

__init__(sigma)

Initialize the AnalyticGaussianNoiseModel.

Parameters:

sigma (float) – Standard deviation of the Gaussian noise.

estimate_conditional_entropy(images=None)

Compute the conditional entropy H(Y | X) for Gaussian noise.

Parameters:

images (jax.Array, optional) – Unused in this model. The parameter is kept for compatibility.

Returns:

The computed conditional entropy, given by the formula: 0.5 * log(2 * pi * e * sigma^2).

Return type:

float

class encoding_information.models.conditional_entropy_models.PoissonNoiseModel

Bases: MeasurementNoiseModel

Poisson noise model for estimating the conditional entropy H(Y | X) from empirical data.

estimate_conditional_entropy(images)

Compute the conditional entropy H(Y | X) for Poisson noise.

Parameters:

images (jax.Array) – A dataset of images, preferably clean. This is used to compute the conditional entropy.

Returns:

The average conditional entropy per pixel, computed using a Gaussian approximation.

Return type:

float

Gaussian Process Models

class encoding_information.models.gaussian_process.StationaryGaussianProcess(images, eigenvalue_floor=0.001, seed=None, verbose=False)

Bases: MeasurementModel

Stationary 2D Gaussian process for single-channel images.

This model learns a stationary covariance matrix and uses it to generate or evaluate samples.

Initialize the StationaryGaussianProcess model and estimate the covariance matrix.

Parameters:
  • images (ndarray) – Input images used to estimate the stationary covariance matrix.

  • eigenvalue_floor (float, optional) – The minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).

  • seed (int, optional) – Random seed for reproducibility.

  • verbose (bool, optional) – Whether to print progress and additional information (default is False).

__init__(images, eigenvalue_floor=0.001, seed=None, verbose=False)

Initialize the StationaryGaussianProcess model and estimate the covariance matrix.

Parameters:
  • images (ndarray) – Input images used to estimate the stationary covariance matrix.

  • eigenvalue_floor (float, optional) – The minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).

  • seed (int, optional) – Random seed for reproducibility.

  • verbose (bool, optional) – Whether to print progress and additional information (default is False).

compute_negative_log_likelihood(images, data_seed=None, verbose=True, seed=None, average=True)

Compute the negative log-likelihood of the provided images under the learned Gaussian process.

Parameters:
  • images (ndarray) – Input images to evaluate.

  • data_seed (int, optional) – Random seed for shuffling the data.

  • verbose (bool, optional) – Whether to print progress (default is True).

  • seed (int, optional, deprecated) – Deprecated argument for random seed, use data_seed instead.

  • average (bool) – If true, return the average NLL, if false return per-example

Returns:

The negative log-likelihood of the provided images.

Return type:

float

fit(train_images=None, data_seed=None, learning_rate=100.0, max_epochs=60, steps_per_epoch=1, patience=15, batch_size=12, num_val_samples=None, percent_samples_for_validation=0.1, eigenvalue_floor=0.001, gradient_clip=1, momentum=0.9, precondition_gradient=False, verbose=True)

Train the Stationary Gaussian Process model on the provided images.

Parameters:
  • train_images (ndarray, optional) – The training dataset consisting of images (default is the images provided during initialization).

  • data_seed (int, optional) – Random seed for shuffling the data.

  • learning_rate (float, optional) – Learning rate for optimization (default is 1e2).

  • max_epochs (int, optional) – Maximum number of training epochs (default is 60).

  • steps_per_epoch (int, optional) – Number of steps per epoch (default is 1).

  • patience (int, optional) – Number of epochs to wait for early stopping (default is 15).

  • batch_size (int, optional) – The number of images in each batch (default is 12).

  • num_val_samples (int, optional) – Number of validation samples (default is computed automatically based on percent_samples_for_validation).

  • percent_samples_for_validation (float, optional) – Fraction of samples to use for validation (default is 0.1).

  • eigenvalue_floor (float, optional) – The minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).

  • gradient_clip (float, optional) – Maximum value to clip the gradient to avoid large updates (default is 1).

  • momentum (float, optional) – Momentum parameter for the optimizer (default is 0.9).

  • precondition_gradient (bool, optional) – Whether to precondition the gradient using the Fisher Information Matrix (default is False).

  • verbose (bool, optional) – Whether to print progress (default is True).

Returns:

val_loss_history – A list of validation losses for each epoch.

Return type:

list

generate_samples(num_samples, sample_shape=None, ensure_nonnegative=True, seed=None, verbose=True)

Generate new image samples from the learned Gaussian process.

Parameters:
  • num_samples (int) – The number of samples to generate.

  • sample_shape (tuple, optional) – Shape of the samples to generate (default is the shape of the training images).

  • ensure_nonnegative (bool, optional) – Whether to ensure all pixel values are non-negative (default is True).

  • seed (int, optional) – Random seed for reproducibility.

  • verbose (bool, optional) – Whether to print progress (default is True).

Returns:

Generated image samples.

Return type:

ndarray

class encoding_information.models.gaussian_process.FullGaussianProcess(data, eigenvalue_floor=0.001, seed=None, verbose=False, add_uniform_noise=True)

Bases: MeasurementModel

Full (non-stationary) Gaussian process for arbitrary data.

This model estimates a full covariance matrix from the data and uses it to generate or evaluate samples.

Initialize the Full Gaussian Process model and estimate the mean and covariance matrix from the data.

Parameters:
  • data (ndarray) – Input data used to estimate the Gaussian process.

  • eigenvalue_floor (float, optional) – Minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).

  • seed (int, optional) – Random seed for reproducibility.

  • verbose (bool, optional) – Whether to print progress during computation (default is False).

  • add_uniform_noise (bool, optional) – For discrete-valued data, add uniform noise (default True). For continuous-valued data, set to False.

__init__(data, eigenvalue_floor=0.001, seed=None, verbose=False, add_uniform_noise=True)

Initialize the Full Gaussian Process model and estimate the mean and covariance matrix from the data.

Parameters:
  • data (ndarray) – Input data used to estimate the Gaussian process.

  • eigenvalue_floor (float, optional) – Minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).

  • seed (int, optional) – Random seed for reproducibility.

  • verbose (bool, optional) – Whether to print progress during computation (default is False).

  • add_uniform_noise (bool, optional) – For discrete-valued data, add uniform noise (default True). For continuous-valued data, set to False.

compute_negative_log_likelihood(data, data_seed=None, verbose=True, seed=None, average=True)

Compute the negative log-likelihood of the provided data under the Gaussian process.

Parameters:
  • data (ndarray) – Input data to evaluate.

  • data_seed (int, optional) – Random seed for shuffling the data.

  • verbose (bool, optional) – Whether to print progress (default is True).

  • seed (int, optional, deprecated) – Deprecated argument for random seed, use data_seed instead.

  • average (bool, optional) – Whether to average the negative log-likelihood over all samples (default is True).

Returns:

The negative log-likelihood of the provided data.

Return type:

float

generate_samples(num_samples, sample_shape=None, ensure_nonnegative=True, seed=None, verbose=True)

Generate new samples from the learned Gaussian process.

Parameters:
  • num_samples (int) – Number of samples to generate.

  • sample_shape (tuple, optional) – Shape of the samples to generate (default is the shape of the training data).

  • ensure_nonnegative (bool, optional) – Whether to ensure all values are non-negative (default is True).

  • seed (int, optional) – Random seed for reproducibility.

  • verbose (bool, optional) – Whether to print progress (default is True).

Returns:

Generated samples.

Return type:

ndarray

Pixel CNN Model

class encoding_information.models.pixel_cnn.PixelCNN(num_hidden_channels=64, num_mixture_components=40)

Bases: MeasurementModel

The PixelCNN model for autoregressive image modeling.

This class handles the training and evaluation of the PixelCNN model and wraps the Flax implementation in a higher-level interface that conforms to the MeasurementModel class. It provides methods for fitting the model to data, computing the negative log-likelihood of images, and generating new images.

num_hidden_channels

The number of hidden channels in the model.

Type:

int

num_mixture_components

The number of components in the mixture density output.

Type:

int

Initialize the PixelCNN model with image shape, number of hidden channels, and mixture components.

Parameters:
  • num_hidden_channels (int) – Number of hidden channels in the convolutional layers.

  • num_mixture_components (int) – Number of mixture components for the output layer.

__init__(num_hidden_channels=64, num_mixture_components=40)

Initialize the PixelCNN model with image shape, number of hidden channels, and mixture components.

Parameters:
  • num_hidden_channels (int) – Number of hidden channels in the convolutional layers.

  • num_mixture_components (int) – Number of mixture components for the output layer.

compute_negative_log_likelihood(data, conditioning_vecs=None, data_seed=None, average=True, verbose=True, seed=None)

Compute the negative log-likelihood (NLL) of images under the trained PixelCNN model.

Parameters:
  • data (ndarray) – The input images for which to compute the NLL.

  • conditioning_vecs (ndarray, optional) – Vectors to condition the image generation process (e.g., class labels).

  • data_seed (int, optional) – Seed for data shuffling.

  • average (bool, optional) – If True, return the average NLL over all images (default is True).

  • verbose (bool, optional) – Whether to print progress (default is True).

  • seed (int, optional) – Deprecated. Use data_seed instead.

Returns:

nll – The negative log-likelihood of the input images.

Return type:

float

fit(train_images, condition_vectors=None, learning_rate=0.01, max_epochs=200, steps_per_epoch=100, patience=40, sigma_min=1, batch_size=64, num_val_samples=None, percent_samples_for_validation=0.1, do_lr_decay=False, verbose=True, add_gaussian_noise=False, add_uniform_noise=True, model_seed=None, data_seed=None, use_positional_embedding=False, seed=None)

Train the PixelCNN model on a dataset of images.

Parameters:
  • train_images (ndarray) – The input dataset, with shape (N, H, W, C).

  • condition_vectors (ndarray, optional) – Vectors to condition the image generation process (e.g., class labels).

  • learning_rate (float, optional) – The learning rate for optimization (default is 1e-2).

  • max_epochs (int, optional) – The maximum number of training epochs (default is 200).

  • steps_per_epoch (int, optional) – The number of steps per epoch (default is 100).

  • patience (int, optional) – The number of epochs to wait before early stopping (default is 40).

  • sigma_min (float, optional) – The minimum standard deviation for the mixture density output (default is 1).

  • batch_size (int, optional) – The batch size for training (default is 64).

  • num_val_samples (int, optional) – The number of validation samples. If None, a percentage is used (default is None).

  • percent_samples_for_validation (float, optional) – The percentage of samples to use for validation (default is 0.1).

  • do_lr_decay (bool, optional) – Whether to apply learning rate decay during training (default is False).

  • verbose (bool, optional) – Whether to print progress during training (default is True).

  • add_gaussian_noise (bool, optional) – Whether to add Gaussian noise to the training images (default is False).

  • add_uniform_noise (bool, optional) – Whether to add uniform noise to the training images (default is True).

  • model_seed (int, optional) – Seed for model initialization.

  • data_seed (int, optional) – Seed for data shuffling.

Returns:

val_loss_history – A list of validation loss values for each epoch.

Return type:

list

generate_samples(num_samples, conditioning_vecs=None, sample_shape=None, ensure_nonnegative=True, seed=None, verbose=True)

Generate new images from the trained PixelCNN model by sampling pixel by pixel.

Parameters:
  • num_samples (int) – Number of images to generate.

  • conditioning_vecs (jax.Array, optional) – Optional conditioning vectors. If provided, the shape should match (num_samples, condition_vector_size). Default is None.

  • sample_shape (tuple of int or int, optional) – Shape of the images to generate. If None, the model’s image_shape is used. If a single int is provided, it will be treated as a square shape. Default is None.

  • ensure_nonnegative (bool, optional) – If True, ensure that the generated pixel values are non-negative. Default is True.

  • seed (int, optional) – Random seed for reproducibility. Default is 123 if not provided.

  • verbose (bool, optional) – If True, display progress during the generation process. Default is True.

Returns:

Generated images with the specified shape.

Return type:

jax.Array